# wzry_core.py - 核心查询类
"""
王者荣耀查询核心模块
优化版本，支持高速查询和多设备管理
"""

import uiautomator2 as u2
import time
import threading
from typing import Dict, List, Optional, Tuple, Any
from loguru import logger
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor


@dataclass
class PlayerResult:
    """玩家查询结果数据类"""
    nickname: str          # 玩家昵称
    role_name: str         # 角色名称  
    server_info: str       # 区服信息
    rank_info: str         # 段位信息
    search_time: float     # 查询耗时


class WZRYDevice:
    """单个设备的查询封装"""
    
    def __init__(self, device_id: str):
        self.device_id = device_id
        self.device = None
        self.is_connected = False
        self.is_busy = False
        self.last_use_time = time.time()
        self.current_page_status = "unknown"
        self.lock = threading.Lock()
        
    def connect(self) -> bool:
        """连接设备"""
        try:
            self.device = u2.connect(self.device_id)
            self.device.implicitly_wait(2.0)
            self.is_connected = True
            self.last_use_time = time.time()
            logger.info(f"设备 {self.device_id} 连接成功")
            return True
        except Exception as e:
            logger.error(f"设备 {self.device_id} 连接失败: {e}")
            self.is_connected = False
            return False
    
    def disconnect(self):
        """断开设备连接"""
        self.is_connected = False
        self.device = None
        logger.info(f"设备 {self.device_id} 已断开")
    
    def is_available(self) -> bool:
        """检查设备是否可用"""
        return self.is_connected and not self.is_busy
    
    def acquire(self) -> bool:
        """获取设备使用权"""
        with self.lock:
            if self.is_available():
                self.is_busy = True
                self.last_use_time = time.time()
                return True
            return False
    
    def release(self):
        """释放设备使用权"""
        with self.lock:
            self.is_busy = False
            self.last_use_time = time.time()


class WZRYDevicePool:
    """设备连接池管理器"""
    
    def __init__(self):
        self.devices: Dict[str, WZRYDevice] = {}
        self.executor = ThreadPoolExecutor(max_workers=10)
        
    def add_device(self, device_id: str) -> bool:
        """添加设备到连接池"""
        if device_id not in self.devices:
            device = WZRYDevice(device_id)
            if device.connect():
                self.devices[device_id] = device
                logger.info(f"设备 {device_id} 已添加到连接池")
                return True
        return False
    
    def remove_device(self, device_id: str):
        """从连接池移除设备"""
        if device_id in self.devices:
            self.devices[device_id].disconnect()
            del self.devices[device_id]
            logger.info(f"设备 {device_id} 已从连接池移除")
    
    def get_available_device(self) -> Optional[WZRYDevice]:
        """获取可用设备"""
        for device in self.devices.values():
            if device.acquire():
                return device
        return None
    
    def get_device_status(self) -> Dict[str, Dict]:
        """获取所有设备状态"""
        status = {}
        for device_id, device in self.devices.items():
            status[device_id] = {
                "connected": device.is_connected,
                "busy": device.is_busy,
                "last_use": device.last_use_time,
                "page_status": device.current_page_status
            }
        return status


class WZRYQueryEngine:
    """王者荣耀查询引擎"""
    
    def __init__(self, device_pool: WZRYDevicePool):
        self.device_pool = device_pool
        
    def _detect_page_status(self, device: u2.Device) -> str:
        """快速检测页面状态"""
        try:
            # 使用更快的方式检测页面
            if device(resourceId="com.tencent.gamehelper.smoba:id/filter_title").exists(timeout=1):
                return "search_result"
            elif device(text="搜索游戏昵称添加好友").exists(timeout=1):
                return "search_input"
            else:
                return "unknown"
        except:
            return "unknown"
    
    def _navigate_to_search_page(self, device: u2.Device) -> bool:
        """快速导航到搜索页面"""
        try:
            page_status = self._detect_page_status(device)
            
            if page_status == "search_result":
                # 从结果页返回搜索页
                back_btn = device(resourceId="com.tencent.gamehelper.smoba:id/back")
                if back_btn.exists(timeout=2):
                    back_btn.click()
                    time.sleep(1)
                    return True
            elif page_status == "search_input":
                # 已在搜索页面
                return True
            
            return False
        except Exception as e:
            logger.error(f"导航到搜索页面失败: {e}")
            return False
    
    def _fast_input_and_search(self, device: u2.Device, player_name: str) -> bool:
        """快速输入和搜索"""
        try:
            # 查找输入框（支持两种类型）
            input_box = None
            for resource_id in ["com.tencent.gamehelper.smoba:id/name_edit", 
                              "com.tencent.gamehelper.smoba:id/search_edit"]:
                box = device(resourceId=resource_id)
                if box.exists(timeout=1):
                    input_box = box
                    break
            
            if not input_box:
                return False
            
            # 快速清空并输入
            input_box.click()
            time.sleep(0.2)
            
            # 尝试清空
            try:
                input_box.clear_text()
                time.sleep(0.1)
            except:
                pass
            
            # 如果清空失败，使用选择删除
            try:
                input_box.long_click()
                time.sleep(0.3)
                if device(text="全选").exists(timeout=1):
                    device(text="全选").click()
                    time.sleep(0.2)
                device.press("del")
            except:
                pass
            
            # 输入搜索内容
            input_box.send_keys(player_name)
            time.sleep(0.3)
            
            # 触发搜索
            search_btn = device(resourceId="com.tencent.gamehelper.smoba:id/search_1")
            if search_btn.exists(timeout=1):
                search_btn.click()
            else:
                device.press("enter")
            
            # 等待搜索结果
            time.sleep(2)
            return True
            
        except Exception as e:
            logger.error(f"输入搜索失败: {e}")
            return False
    
    def _parse_search_results(self, device: u2.Device, search_name: str) -> List[PlayerResult]:
        """解析搜索结果"""
        try:
            results = []
            
            # 获取所有结果元素
            nickname_elements = device(resourceId="com.tencent.gamehelper.smoba:id/nickname")
            role_name_elements = device(resourceId="com.tencent.gamehelper.smoba:id/role_name")
            server_info_elements = device(resourceId="com.tencent.gamehelper.smoba:id/server_info")
            rank_info_elements = device(resourceId="com.tencent.gamehelper.smoba:id/jon_info")
            
            # 解析每个结果
            count = len(nickname_elements)
            for i in range(count):
                try:
                    nickname = nickname_elements[i].get_text() if i < len(nickname_elements) else "未知"
                    role_name = role_name_elements[i].get_text() if i < len(role_name_elements) else "未知"
                    server_info = server_info_elements[i].get_text() if i < len(server_info_elements) else "未知"
                    rank_info = rank_info_elements[i].get_text() if i < len(rank_info_elements) else "未知"
                    
                    # 检查是否匹配搜索条件
                    if self._is_match(search_name, role_name, nickname):
                        result = PlayerResult(
                            nickname=nickname,
                            role_name=role_name,
                            server_info=server_info,
                            rank_info=rank_info,
                            search_time=time.time()
                        )
                        results.append(result)
                        
                except Exception as e:
                    logger.warning(f"解析第{i+1}个结果失败: {e}")
                    continue
            
            return results
            
        except Exception as e:
            logger.error(f"解析搜索结果失败: {e}")
            return []
    
    def _is_match(self, search_name: str, role_name: str, nickname: str) -> bool:
        """检查是否匹配搜索条件"""
        # 支持部分匹配
        search_parts = search_name.split()
        for part in search_parts:
            if part in role_name or part in nickname:
                return True
        return search_name in role_name or search_name in nickname
    
    def search_player(self, player_name: str, timeout: int = 30) -> Tuple[List[PlayerResult], str]:
        """搜索玩家（高性能版本）"""
        start_time = time.time()
        
        # 获取可用设备
        device_obj = self.device_pool.get_available_device()
        if not device_obj:
            return [], "没有可用设备"
        
        device_id = device_obj.device_id
        device = device_obj.device
        
        try:
            logger.info(f"使用设备 {device_id} 搜索玩家: {player_name}")
            
            # 1. 快速导航到搜索页面
            if not self._navigate_to_search_page(device):
                return [], f"设备 {device_id} 导航到搜索页面失败"
            
            # 2. 快速输入和搜索
            if not self._fast_input_and_search(device, player_name):
                return [], f"设备 {device_id} 搜索执行失败"
            
            # 3. 解析结果
            results = self._parse_search_results(device, player_name)
            
            # 4. 记录查询耗时
            elapsed_time = time.time() - start_time
            logger.info(f"设备 {device_id} 查询完成，耗时 {elapsed_time:.2f}秒，找到 {len(results)} 个结果")
            
            return results, "查询成功"
            
        except Exception as e:
            logger.error(f"设备 {device_id} 查询异常: {e}")
            return [], f"设备 {device_id} 查询异常: {str(e)}"
        
        finally:
            # 释放设备
            device_obj.release()


# 全局实例
device_pool = WZRYDevicePool()
query_engine = WZRYQueryEngine(device_pool)